import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt
from matplotlib import rcParams
import numpy as np

# 设置支持中文的字体
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['SimHei']  # 使用中文字体

def rotate_image(image, angle):
    """旋转图像并保持尺寸"""
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)

    # 获得旋转矩阵，然后应用仿射变换
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(image, M, (w, h))
    return rotated

def infer(image, model, angle):
    """对输入图像在指定角度下进行推理"""
    # 对图像进行旋转
    rotated_img = rotate_image(image, angle)

    # 转换为RGB模式（YOLO要求）
    img_rgb = cv2.cvtColor(rotated_img, cv2.COLOR_BGR2RGB)

    # 执行推理
    results = model(img_rgb)

    # 可视化检测结果
    result_img = results[0].plot()  # 绘制检测结果
    plt.imshow(result_img)
    plt.title(f'推理结果 (旋转角度: {angle}度)')
    plt.axis('off')
    plt.show()

    # 输出检测的类别、置信度和边界框信息
    print(f"旋转角度: {angle}度")
    for detection in results[0].boxes.data:
        x1, y1, x2, y2, score, class_id = detection.tolist()
        print(f"类别: {model.names[int(class_id)]}, 置信度: {score:.2f}, 边界框: ({x1}, {y1}, {x2}, {y2})")

def main():
    # 加载训练好的YOLO模型
    model = YOLO('./best.pt')  # 使用训练后保存的模型

    # 加载图像
    image_path = './image.jpg'
    img = cv2.imread(image_path)

    # 对不同旋转角度的图像进行推理
    angles = [0, 10, 20, 30, 45, 90]  # 定义旋转角度
    for angle in angles:
        infer(img, model, angle)

if __name__ == '__main__':
    main()
